package kqbuffer

import (
	"errors"
	"sync/atomic"

	"github.com/segmentio/kafka-go"
)

var (
	ErrEmptyMsg     = errors.New("empty kafka Message")
	ErrResetNoEmpty = errors.New("reset no empty buffer")
	ErrPushFull     = errors.New("push full kqutil buffer")
	ErrPopEmpty     = errors.New("pop empty kqutil buffer")
	ErrPopPending   = errors.New("pop kqutil buffer pending")
)

type (
	kqBufferRange struct {
		start int
		end   int
	}

	KqBatchMsg struct {
		Msgs []*kafka.Message
		r    kqBufferRange
	}

	PartitionOffsetMap map[int]*kafka.Message
)

/*
 * 特点：
 * 作为 write buffer 时只进不出，作为 read buffer 只出不进
 */
type (
	KqWriteBuffer interface {
		AsRead() KqReadBuffer
		Push(msg *kafka.Message) error
		IsFull() bool
		IsEmpty() bool
		Length() int
	}

	KqReadBuffer interface {
		IsBatchEmpty() bool
		Batch(size int) (*KqBatchMsg, error)
		Ready(batchMsg *KqBatchMsg) (countWaitReady int)
		Pop() (kqOffset PartitionOffsetMap, err error)
		IsFull() bool
		IsEmpty() bool
		Length() int
		ResetBuffer() error
	}

	kqBuffer struct {
		msgBuffer   []*kafka.Message // kafka消息缓冲区队列
		readyCommit []bool           // 与消息列表位置一一对应，表示当前消息是否处理完毕，可以提交

		head     int
		tail     int
		capacity int

		// read buffer 特供
		readIndex      int
		countWaitReady atomic.Int32
	}
)

func newKqBuffer(capacity int) *kqBuffer {
	return &kqBuffer{
		msgBuffer:   make([]*kafka.Message, capacity),
		readyCommit: make([]bool, capacity),

		head:     0,
		tail:     0,
		capacity: capacity,

		readIndex: 0,
	}
}

/*
 * 添加消息，返回添加结果
 */
func (b *kqBuffer) Push(msg *kafka.Message) error {
	if msg == nil {
		return ErrEmptyMsg
	}
	if b.IsFull() {
		return ErrPushFull
	}
	b.msgBuffer[b.tail] = msg
	b.readyCommit[b.tail] = false
	b.tail++
	return nil
}

/*
 * 转换为读缓存
 */
func (b *kqBuffer) AsRead() KqReadBuffer {
	b.countWaitReady.Store(int32(b.Length()))
	return b
}

/*
 * 判断当前是否还有剩余未分发 batch msg
 */
func (b *kqBuffer) IsBatchEmpty() bool {
	return b.readIndex == b.tail
}

/*
 * 获取最多 size 数量的结果，返回获取结果
 */
func (b *kqBuffer) Batch(size int) (*KqBatchMsg, error) {
	start := b.readIndex
	end := min(start+size, b.tail)
	if start == end {
		return nil, ErrEmptyMsg
	}
	b.readIndex = end
	return &KqBatchMsg{
		Msgs: b.msgBuffer[start:end],
		r:    kqBufferRange{start, end},
	}, nil
}

/*
 * 标记 batchMsg 已消费，返回剩余 Ready 数量
 */
func (b *kqBuffer) Ready(batchMsg *KqBatchMsg) (countWaitReady int) {
	if batchMsg == nil {
		return int(b.countWaitReady.Load())
	}
	if batchMsg.r.start < b.head || batchMsg.r.start >= b.tail ||
		batchMsg.r.end < b.head || batchMsg.r.end > b.tail {
		panic("TrackRange out of bounds")
	}

	var readyCount int32 = 0
	for i := batchMsg.r.start; i < batchMsg.r.end; i++ {
		if !b.readyCommit[i] {
			b.readyCommit[i] = true
			readyCount++
		}
	}
	return int(b.countWaitReady.Add(-readyCount))
}

/*
 * 清理已完成的消费，返回当前缓冲区各分区已完成消费的最后一条消息
 */
func (d *kqBuffer) Pop() (kqOffset PartitionOffsetMap, err error) {
	if d.IsEmpty() {
		return nil, ErrPopEmpty
	}

	kqOffset = make(PartitionOffsetMap)
	for d.head < d.tail {
		if !d.readyCommit[d.head] {
			break
		}
		msg := d.msgBuffer[d.head]
		kqOffset[msg.Partition] = msg
		d.readyCommit[d.head] = false
		d.msgBuffer[d.head] = nil
		d.head++
	}

	if len(kqOffset) == 0 {
		return nil, ErrPopPending
	}

	return
}

func (b *kqBuffer) IsFull() bool {
	return b.Length() == b.capacity
}

func (b *kqBuffer) IsEmpty() bool {
	return b.Length() == 0
}

func (b *kqBuffer) Length() int {
	return b.tail - b.head
}

/*
 * read buffer 特供，因为只有读缓冲区能清空
 */
func (b *kqBuffer) ResetBuffer() error {
	if !b.IsEmpty() {
		return ErrResetNoEmpty
	}
	b.head = 0
	b.tail = 0
	b.readIndex = 0
	return nil
}
